# -*- coding: utf-8 -*-
"""
GeoJSON文件处理器
将GeoJSON文件转换为标准化的PostGIS导入参数格式
"""

import geopandas as gpd
import pandas as pd
import json
from typing import Dict, Any, Optional
import os
from loguru import logger


class GeoJSONProcessor:
    """GeoJSON文件处理器，转换为PostGIS导入格式"""
    
    def __init__(self, file_path: str):
        self.file_path = file_path
        self.gdf = None
        
    def load_data(self) -> Dict[str, Any]:
        """
        加载GeoJSON文件并转换为标准格式
        
        Returns:
            Dict包含: {
                'data': 标准化后的GeoDataFrame对应表结构最后存到PostGIS中的数据,
                'geometry_type': 几何类型,
                'crs': 坐标系,
                'columns': 字段列表,
                'record_count': 记录数量,
                'feature_count': 原始特征数量
            }
        """
        try:
            # 读取GeoJSON文件
            self.gdf = gpd.read_file(self.file_path)
            
            # 标准化数据结构
            standardized_data = self._standardize_data()
            
            result = {
                'data': standardized_data,
                'geometry_type': self._get_geometry_type(),
                'crs': str(self.gdf.crs) if self.gdf.crs else 'EPSG:4326',
                'columns': list(standardized_data.columns),
                'record_count': len(standardized_data),
                'feature_count': len(self.gdf),
                'source_format': 'GeoJSON'
            }
            
            logger.info(f"GeoJSON处理完成: {os.path.basename(self.file_path)}")
            logger.info(f"特征数: {result['feature_count']}, 几何类型: {result['geometry_type']}")
            
            return result
            
        except Exception as e:
            logger.error(f"GeoJSON文件处理失败: {e}")
            raise
    
    def _standardize_data(self) -> gpd.GeoDataFrame:
        """标准化数据结构"""
        if self.gdf is None:
            raise ValueError("请先加载数据")
            
        # 创建标准化列
        standardized = self.gdf.copy()
        
        # 处理嵌套属性
        if 'properties' in standardized.columns:
            # 展开properties字段
            properties_df = pd.json_normalize(standardized['properties'].apply(
                lambda x: x if isinstance(x, dict) else {}
            ))
            standardized = pd.concat([standardized.drop('properties', axis=1), properties_df], axis=1)
        
        # 确保有标准字段
        if 'code' not in standardized.columns:
            # 优先尝试stcd字段作为code，然后尝试其他常见水文字段
            code_candidates = ['stcd', 'STCD', 'stcd_field', 'gd', 'GD', 'id', 'ID', 'code', 'Code', 'CODE', 'station_id', 'StationID']
            for candidate in code_candidates:
                if candidate in standardized.columns:
                    standardized['code'] = standardized[candidate].astype(str)
                    logger.info(f"使用字段 '{candidate}' 作为code")
                    break
            else:
                standardized['code'] = [f"GEO_{i:06d}" for i in range(len(standardized))]
                logger.warning("未找到合适的code字段，使用自动生成")

        if 'name' not in standardized.columns:
            # 尝试从更多可能的名称字段获取
            name_candidates = ['name', 'Name', 'NAME', 'stnm', 'STNM', 'title', 'label', 'station_name', 'reservoir_name', 'lake_name']
            for candidate in name_candidates:
                if candidate in standardized.columns:
                    standardized['name'] = standardized[candidate].astype(str)
                    logger.info(f"使用字段 '{candidate}' 作为name")
                    break
            else:
                # 如果没有合适的名称字段，尝试使用code作为name
                if 'code' in standardized.columns:
                    standardized['name'] = standardized['code']
                    logger.info("使用code作为name")
                else:
                    standardized['name'] = [f"未命名_{i}" for i in range(len(standardized))]
                
        if 'type' not in standardized.columns:
            # 尝试从类型相关字段获取
            type_candidates = ['type', 'Type', 'TYPE', 'category', 'Category', 'kind', 'Kind', 'reservoir_type', 'station_type']
            for candidate in type_candidates:
                if candidate in standardized.columns:
                    standardized['type'] = standardized[candidate].astype(str)
                    logger.info(f"使用字段 '{candidate}' 作为type")
                    break
            else:
                # 根据几何类型推断合理的默认值
                first_geom_type = standardized.geometry.iloc[0].geom_type if len(standardized) > 0 else 'Unknown'
                if first_geom_type in ['Point', 'MultiPoint']:
                    standardized['type'] = '水文站点'
                elif first_geom_type in ['Polygon', 'MultiPolygon']:
                    standardized['type'] = '水库'
                else:
                    standardized['type'] = '要素'
        else:
            # 处理空值、null值和空字符串
            mask = standardized['type'].isna() | (standardized['type'] == '') | (standardized['type'].astype(str).str.strip() == '')
            standardized.loc[mask, 'type'] = '要素'
            
        # 多余字段直接丢弃，extra_data保持为空
        standardized['extra_data'] = ''
            
        # 转换几何类型以匹配数据库要求
        standardized = self._convert_geometry_types(standardized)
        
        # 转换坐标系为EPSG:4326
        if standardized.crs is None:
            standardized = standardized.set_crs('EPSG:4326')
        elif str(standardized.crs) != 'EPSG:4326':
            standardized = standardized.to_crs('EPSG:4326')
            
        # 只保留标准字段
        result = standardized[['code', 'name', 'type', 'extra_data', 'geometry']].copy()
        
        return result
    
    def _convert_geometry_types(self, gdf: gpd.GeoDataFrame) -> gpd.GeoDataFrame:
        """转换几何类型以匹配数据库要求"""
        if gdf is None or len(gdf) == 0:
            return gdf
            
        # 创建副本避免修改原始数据
        result = gdf.copy()
        
        # 获取原始几何类型
        first_geom_type = result.geometry.iloc[0].geom_type
        
        # 根据类型进行转换
        from shapely.geometry import MultiPolygon, Point
        
        if first_geom_type == 'Polygon':
            # 将Polygon转换为MultiPolygon
            result.geometry = result.geometry.apply(
                lambda geom: geom if geom is None else 
                geom if geom.geom_type == 'MultiPolygon' else
                geom if geom.geom_type != 'Polygon' else
                MultiPolygon([geom])
            )
        elif first_geom_type == 'MultiPoint':
            # 将MultiPoint转换为Point（取第一个点）
            result.geometry = result.geometry.apply(
                lambda geom: geom if geom is None else
                geom if geom.geom_type == 'Point' else
                geom if geom.geom_type != 'MultiPoint' else
                list(geom.geoms)[0] if len(list(geom.geoms)) > 0 else None
            )
            
        return result
    
    def _get_geometry_type(self) -> str:
        """获取几何类型"""
        if self.gdf is None or len(self.gdf) == 0:
            return 'Unknown'
            
        geom_type = self.gdf.geometry.iloc[0].geom_type
        
        # 标准化几何类型
        type_mapping = {
            'Polygon': 'MULTIPOLYGON',
            'MultiPolygon': 'MULTIPOLYGON',
            'Point': 'POINT',
            'MultiPoint': 'POINT'
        }
        
        return type_mapping.get(geom_type, geom_type.upper())
    
    def validate_data(self) -> Dict[str, Any]:
        """验证数据完整性"""
        if self.gdf is None:
            return {'valid': False, 'errors': ['数据未加载']}
            
        errors = []
        warnings = []
        
        # 检查几何有效性
        invalid_geoms = ~self.gdf.geometry.is_valid
        if invalid_geoms.any():
            errors.append(f"发现{invalid_geoms.sum()}个无效几何图形")
            
        # 检查空几何
        empty_geoms = self.gdf.geometry.is_empty
        if empty_geoms.any():
            errors.append(f"发现{empty_geoms.sum()}个空几何图形")
            
        # 检查重复code
        if 'code' in self.gdf.columns:
            duplicate_codes = self.gdf['code'].duplicated().sum()
            if duplicate_codes > 0:
                errors.append(f"发现{duplicate_codes}个重复code")
                
        # 检查坐标系
        if self.gdf.crs is None:
            warnings.append("未定义坐标系，将使用EPSG:4326")
            
        return {
            'valid': len(errors) == 0,
            'errors': errors,
            'warnings': warnings
        }
    
    def get_metadata(self) -> Dict[str, Any]:
        """获取文件元数据"""
        if self.gdf is None:
            return {}
            
        return {
            'bbox': self.gdf.total_bounds.tolist(),
            'crs': str(self.gdf.crs),
            'geometry_types': self.gdf.geometry.geom_type.value_counts().to_dict(),
            'field_types': self.gdf.dtypes.to_dict()
        }


# 测试示例
if __name__ == "__main__":
    import tempfile
    
    # 创建临时测试文件
    test_file = "test_data.geojson"
    
    # 创建测试GeoJSON数据
    test_geojson = {
        "type": "FeatureCollection",
        "features": [
            {
                "type": "Feature",
                "properties": {
                    "name": "测试水库1",
                    "capacity": 1000,
                    "type": "大型水库"
                },
                "geometry": {
                    "type": "Point",
                    "coordinates": [116.4074, 39.9042]
                }
            },
            {
                "type": "Feature", 
                "properties": {
                    "name": "测试水库2",
                    "capacity": 500,
                    "type": "中型水库"
                },
                "geometry": {
                    "type": "Point",
                    "coordinates": [117.1993, 39.0851]
                }
            },
            {
                "type": "Feature",
                "properties": {
                    "name": "测试MultiPoints",
                    "capacity": 800,
                    "type": "多点测试"
                },
                "geometry": {
                    "type": "MultiPoint",
                    "coordinates": [
                        [116.4074, 39.9042],
                        [117.1993, 39.0851],
                        [118.1953, 39.6321]
                    ]
                }
            },
            {
                "type": "Feature",
                "properties": {
                    "name": "测试Polygon",
                    "capacity": 1200,
                    "type": "多边形测试"
                },
                "geometry": {
                    "type": "Polygon",
                    "coordinates": [
                        [
                            [116.0, 39.0],
                            [117.0, 39.0],
                            [117.0, 40.0],
                            [116.0, 40.0],
                            [116.0, 39.0]
                        ]
                    ]
                }
            }
        ]
    }
    
    # 保存测试数据到文件
    with open(test_file, 'w', encoding='utf-8') as f:
        json.dump(test_geojson, f, ensure_ascii=False, indent=2)
    
    try:
        processor = GeoJSONProcessor(test_file)
        
        # 加载并处理数据
        result = processor.load_data()
        
        # 验证数据
        validation = processor.validate_data()
        
        print("=== GeoJSON处理结果 ===")
        print(f"文件: {test_file}")
        print(f"记录数: {result['record_count']}")
        print(f"几何类型: {result['geometry_type']}")
        print(f"坐标系: {result['crs']}")
        print(f"字段: {result['columns']}")
        print(f"数据验证: {'通过' if validation['valid'] else '失败'}")
        
        if validation['warnings']:
            print("警告:", validation['warnings'])
        if not validation['valid']:
            print("错误:", validation['errors'])
            
        # 显示元数据
        metadata = processor.get_metadata()
        print("\n=== 元数据 ===")
        print(f"边界框: {metadata.get('bbox', 'N/A')}")
        print(f"几何类型分布: {metadata.get('geometry_types', {})}")
        
        # 显示数据示例
        print("\n=== 数据示例 ===")
        print(result['data'][['code', 'name', 'type', 'extra_data']].head())
        
    finally:
        # 清理测试文件
        if os.path.exists(test_file):
            os.remove(test_file)
    
    # 本地文件测试示例
    print("\n=== 本地文件测试 ===")
    local_file_path = "碧流河.json"  # 请在这里填写本地GeoJSON文件路径，例如: "碧流河.json"
    
    if local_file_path and os.path.exists(local_file_path):
        try:
            processor = GeoJSONProcessor(local_file_path)
            
            # 加载并处理数据
            result = processor.load_data()
            
            # 验证数据
            validation = processor.validate_data()
            
            print(f"=== 本地文件处理结果: {local_file_path} ===")
            print(f"记录数: {result['record_count']}")
            print(f"几何类型: {result['geometry_type']}")
            print(f"坐标系: {result['crs']}")
            print(f"字段: {result['columns']}")
            print(f"数据验证: {'通过' if validation['valid'] else '失败'}")
            
            if validation['warnings']:
                print("警告:", validation['warnings'])
            if not validation['valid']:
                print("错误:", validation['errors'])
                
            # 显示数据示例
            print("\n=== 数据示例 ===")
            print(result['data'][['code', 'name', 'type', 'extra_data']].head())
            
        except Exception as e:
            print(f"处理本地文件失败: {e}")
    else:
        print("请填写本地GeoJSON文件路径进行测试")